{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b3e78a24",
      "metadata": {},
      "outputs": [],
      "source": [
        "from __future__ import annotations\n",
        "\n",
        "import io\n",
        "from typing import Any, Dict, List\n",
        "\n",
        "import av\n",
        "import torch\n",
        "import torchvision\n",
        "from PIL import Image\n",
        "from pyspark.sql.functions import col, explode, pandas_udf\n",
        "from pyspark.sql.types import (\n",
        "    ArrayType,\n",
        "    BinaryType,\n",
        "    FloatType,\n",
        "    IntegerType,\n",
        "    StringType,\n",
        "    StructField,\n",
        "    StructType,\n",
        ")\n",
        "from ultralytics import YOLO"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ffe55226",
      "metadata": {},
      "outputs": [],
      "source": [
        "%%configure -f\n",
        "{\n",
        "  \"executorCores\": 1,\n",
        "  \"conf\": {\n",
        "    \"spark.sql.execution.arrow.maxRecordsPerBatch\": \"64\"\n",
        "  }\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb5faa00",
      "metadata": {},
      "outputs": [],
      "source": [
        "_model_cache = {\"model\": None}\n",
        "\n",
        "\n",
        "def get_model():\n",
        "    if _model_cache[\"model\"] is None:\n",
        "        model = YOLO(\"yolo11n.pt\")\n",
        "        if torch.cuda.is_available():\n",
        "            model.to(\"cuda\")\n",
        "        _model_cache[\"model\"] = model\n",
        "    return _model_cache[\"model\"]\n",
        "\n",
        "\n",
        "@pandas_udf(ArrayType(BinaryType()))\n",
        "def decode_video_udf(video_bytes_iter):\n",
        "    import pandas as pd\n",
        "\n",
        "    results = []\n",
        "    for video_bytes in video_bytes_iter:\n",
        "        frames = []\n",
        "        with av.open(io.BytesIO(video_bytes)) as container:\n",
        "            for frame in container.decode(video=0):\n",
        "                img = frame.to_ndarray(format=\"rgb24\")\n",
        "                pil_img = Image.fromarray(img).resize((640, 640))\n",
        "                buf = io.BytesIO()\n",
        "                pil_img.save(buf, format=\"PNG\")\n",
        "                frames.append(buf.getvalue())\n",
        "        results.append(frames)\n",
        "    return pd.Series(results)\n",
        "\n",
        "\n",
        "feature_schema = ArrayType(\n",
        "    StructType(\n",
        "        [\n",
        "            StructField(\"label\", StringType(), False),\n",
        "            StructField(\"confidence\", FloatType(), False),\n",
        "            StructField(\"bbox\", ArrayType(IntegerType()), False),\n",
        "        ]\n",
        "    )\n",
        ")\n",
        "\n",
        "\n",
        "def to_features(res: Any) -> List[Dict[str, Any]]:\n",
        "    return [\n",
        "        {\n",
        "            \"label\": res.names[int(cls_id)],\n",
        "            \"confidence\": float(conf),\n",
        "            \"bbox\": [int(v) for v in bbox.tolist()],\n",
        "        }\n",
        "        for cls_id, conf, bbox in zip(res.boxes.cls, res.boxes.conf, res.boxes.xyxy)\n",
        "    ]\n",
        "\n",
        "\n",
        "@pandas_udf(feature_schema)\n",
        "def extract_image_features_udf(images):\n",
        "    import pandas as pd\n",
        "\n",
        "    if len(images) == 0:\n",
        "        return []\n",
        "    model = get_model()\n",
        "    tensors = [torchvision.transforms.functional.to_tensor(Image.open(io.BytesIO(img))) for img in images]\n",
        "    stack = torch.stack(tensors, dim=0)\n",
        "    results = model(stack)\n",
        "    return pd.Series([to_features(r) for r in results])\n",
        "\n",
        "\n",
        "@pandas_udf(BinaryType())\n",
        "def crop_udf(frame_bytes_iter, bbox_iter):\n",
        "    import pandas as pd\n",
        "\n",
        "    outputs = []\n",
        "    for frame_bytes, bbox in zip(frame_bytes_iter, bbox_iter):\n",
        "        try:\n",
        "            img = Image.open(io.BytesIO(frame_bytes)).convert(\"RGB\")\n",
        "            x1, y1, x2, y2 = bbox\n",
        "            cropped = img.crop((x1, y1, x2, y2))\n",
        "            buf = io.BytesIO()\n",
        "            cropped.save(buf, format=\"PNG\")\n",
        "            outputs.append(buf.getvalue())\n",
        "        except Exception:\n",
        "            outputs.append(None)\n",
        "    return pd.Series(outputs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "54773d95",
      "metadata": {},
      "outputs": [],
      "source": [
        "df = spark.read.format(\"binaryFile\").load(\"s3://daft-public-data/videos/Hollywood2-actions-videos/Hollywood2/AVIClips/\")\n",
        "df = df.withColumn(\"frame\", decode_video_udf(col(\"content\")))\n",
        "df = df.withColumn(\"frame\", explode(col(\"frame\")))\n",
        "df = df.checkpoint()\n",
        "df = df.withColumn(\"features\", extract_image_features_udf(col(\"frame\")))\n",
        "df = df.withColumn(\"feature\", explode(col(\"features\")))\n",
        "df = df.checkpoint()\n",
        "df = df.withColumn(\"object\", crop_udf(col(\"frame\"), col(\"feature.bbox\")))\n",
        "df = df.drop(\"content\", \"frame\")\n",
        "df.write.mode(\"append\").parquet(\"s3://eventual-dev-benchmarking-results/ai-benchmark-results/video-object-detection-result\")"
      ]
    }
  ],
  "metadata": {
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
